import pandas as pd
#from vk_download import main
#import scipy.sparse as sprs
import time
from tqdm.auto import tqdm
import os
import random
import datetime
#from threading import Thread
#from multiprocessing import Process
import numpy as np
from scipy.sparse import csr_matrix, lil_matrix, csc_matrix
#import scipy.sparse as sprs
#from collections import OrderedDict
#import shutil
import matplotlib.pyplot as plt
#import seaborn as sns
import networkx as nx
from IPython.display import clear_output
#from statsmodels.stats.proportion import proportion_confint
#import statsmodels.api as sm
#from scipy.optimize import minimize
import matplotlib as mpl
import copy
from scipy.stats import powerlaw
import warnings
#warnings.filterwarnings('ignore')
#warnings.filterwarnings('always')




################################################################################################################




def UpdateOpinion(Arguments, Records, t, Agent, ElementId, ElementType):
    
    """
    This function models how the agent updates its opinion following the interaction with the news feed element
    """
    
    #------------------------------------ Get the characteristics of the agent and element
    
    TransitionTable = Arguments['TransitionTable']
    
    AgentOpinion = Records['Agents'][Records['Agents'].columns[-1]].loc[Agent]
    
    ElementOpinion = Records[ElementType]['Opinion'].loc[ElementId]
        
    ElementLikesCount = Records[ElementType]['LikesCount'].loc[ElementId]
        
    ElementRepliesCount = Records[ElementType]['RepliesCount'].loc[ElementId]

    #------------------------------------ Update the agent's opinion according to the transition table
    
    AgentNewOpinion = np.random.choice(Arguments['OpinionAlphabet'], p=TransitionTable[AgentOpinion][ElementOpinion])
    
    #------------------------------------ Insert records
    
    NewColumnId = Records['Agents'].shape[1]
    
    Records['Agents'][f'Opinion({t}-{NewColumnId})'] = Records['Agents'][Records['Agents'].columns[-1]]
    
    Records['Agents'].loc[Agent][f'Opinion({t}-{NewColumnId})'] = AgentNewOpinion
    
    #------------------------------------ Update fractions
    
    Records['Fractions'].append(copy.copy(Records['Fractions'][-1]))
    
    #Records['Fractions'].append(Records['Fractions'][-1])
    
    #print(Records['Fractions'][-1])
    
    Records['Fractions'][-1][AgentOpinion] = Records['Fractions'][-1][AgentOpinion] - 1
    
    Records['Fractions'][-1][AgentNewOpinion] = Records['Fractions'][-1][AgentNewOpinion] + 1
    
    #print(Records['Fractions'][-1])
    
    #print('')
    
    return Records




def SocialContagionFunction(Arguments, Counter, Type):
    
    if Type == 'Like':
    
        SocialFactor = Arguments['ScaleFactor_Like'] * (Arguments['Ratio_Like']**Counter - 1) / (Arguments['Ratio_Like'] - 1)
    
    elif Type == 'Reply':
        
        SocialFactor = Arguments['ScaleFactor_Reply'] * (Arguments['Ratio_Reply']**Counter - 1) / (Arguments['Ratio_Reply'] - 1)
        
    else:
        
        raise ValueError('Unknown type of the social contagion object!')
    
    return SocialFactor




def PutALike(Arguments, Records, t, Agent, ElementId, ElementType):
    
    """
    This function models how the agent puts (or does not put) a like on the news feed element
    """
    
    if ElementId not in Records['AgentsLikes'][Agent]:  # check if the like already exists
        
        #------------------------------------ Get the characteristics of the agent and element

        AgentOpinion = Records['Agents'][Records['Agents'].columns[-1]].loc[Agent]  # need the current opinion of the agent

        ElementOpinion = Records[ElementType]['Opinion'].loc[ElementId]

        ElementLikesCount = Records[ElementType]['LikesCount'].loc[ElementId]

        ElementRepliesCount = Records[ElementType]['RepliesCount'].loc[ElementId]
        
        if Agent < Arguments['N_Bots']:  # this agent is a bot
            
            LikingProbabilities = Arguments['LikingProbabilities_Bots']
            
            Probability = LikingProbabilities[AgentOpinion, ElementOpinion]
            
        else:  # this agent is not a bot
            
            LikingProbabilities = Arguments['LikingProbabilities']
            
            #------------------------------------ Compute the social factor (Social Contagion theory)

            BasicProbability = LikingProbabilities[AgentOpinion, ElementOpinion]

            SocialFactor = SocialContagionFunction(Arguments, ElementLikesCount, 'Like')

            Probability = min(BasicProbability + SocialFactor, 1)
        
        #------------------------------------ 

        LikeOrNot = np.random.choice(['Put a like', 'Ignore'], 
                                      p=[Probability, 1-Probability])

        if LikeOrNot == 'Put a like':

            Records['Likes'].loc[Records['Likes'].shape[0]] = [
                                                               t,                       # time creation t=0 is assumed 
                                                               Agent,                   # author of the like
                                                               ElementId,               # object of the like
                                                               ElementType,             # type of the object
                                                              ] 

            #print(ElementType, Records[ElementType].loc[ElementId]['LikesCount'])
            
            #Records['Comments'].loc[0, 'LikesCount'] = 101

            Records[ElementType].loc[ElementId, 'LikesCount'] = ElementLikesCount + 1

            #print(Records[ElementType].loc[ElementId]['LikesCount'])
            
            Records['AgentsLikes'][Agent].append(ElementId)

        else:

            pass
    
    else:
        
        pass
    
    return Records




def WriteReply(Arguments, Records, t, Agent, ElementId, ElementType):
    
    """
    This function models how the agent replies (or does not reply) to the news feed element
    """
    
    #------------------------------------ Get the characteristics of the agent and element
    
    
    
    AgentOpinion = Records['Agents'][Records['Agents'].columns[-1]].loc[Agent]  # need the current opinion of the agent
    
    ElementOpinion = Records[ElementType]['Opinion'].loc[ElementId]
        
    ElementLikesCount = Records[ElementType]['LikesCount'].loc[ElementId]
        
    ElementRepliesCount = Records[ElementType]['RepliesCount'].loc[ElementId]
    
    if Agent < Arguments['N_Bots']:  # this agent is a bot
            
        ReplyWritingProbabilities = Arguments['ReplyWritingProbabilities_Bots']
            
        Probability = ReplyWritingProbabilities[AgentOpinion, ElementOpinion]
            
    else:  # this agent is not a bot
            
        ReplyWritingProbabilities = Arguments['ReplyWritingProbabilities']
            
        #------------------------------------ Compute the social factor (Social Contagion theory)

        BasicProbability = ReplyWritingProbabilities[AgentOpinion, ElementOpinion]

        SocialFactor = SocialContagionFunction(Arguments, ElementRepliesCount, 'Reply')

        Probability = min(BasicProbability + SocialFactor, 1)
    
    #------------------------------------ 
    
    WriteOrNot = np.random.choice(['Write reply', 'Ignore'], 
                                  p=[Probability, 1-Probability])
    
    if WriteOrNot == 'Write reply':
        
        Records['Comments'].loc[Records['Comments'].shape[0]] = [
                                                                 t,                       # time creation t=0 is assumed 
                                                                 AgentOpinion,            # opinion of the reply is defined as the agent's opinion
                                                                 Agent,                   # author of the reply
                                                                 ElementId,               # object of the reply
                                                                 ElementType,             # type of the object
                                                                 0,                       # no likes at the beginning
                                                                 0                        # no replies at the beginning
                                                                ] 
        
        #Records[ElementType]['RepliesCount'].loc[ElementId] = ElementRepliesCount + 1
        
        Records[ElementType].loc[ElementId, 'RepliesCount'] = ElementRepliesCount + 1
        
        #Records['Comments'].loc[0, 'LikesCount'] = 101
    
    else:
        
        pass
    
    return Records




def AgentProceedElement(Arguments, Records, t, Agent, ElementId, ElementType):
    
    """
    This function models how the agent processes a news feed element
    """
    
    # The agent first updates its opinion (if he/she is not a bot)
    
    if Agent < Arguments['N_Bots']:  # this agent is a bot
        
        pass  # bots do no change their opinions
    
    else:
        
        Records = UpdateOpinion(Arguments, Records, t, Agent, ElementId, ElementType)
    
    # Then the agent reacts on the element
    
    if Arguments['OrderOfActions'] == 'LikeFirst':
    
        # The agent puts (or does not put) a like on the element

        Records = PutALike(Arguments, Records, t, Agent, ElementId, ElementType)

        # The agent writes (or does not write) a reply

        Records = WriteReply(Arguments, Records, t, Agent, ElementId, ElementType)
        
    elif Arguments['OrderOfActions'] == 'ReplyFirst':

        # The agent writes (or does not write) a reply

        Records = WriteReply(Arguments, Records, t, Agent, ElementId, ElementType)
        
        # The agent puts (or does not put) a like on the element

        Records = PutALike(Arguments, Records, t, Agent, ElementId, ElementType)
        
    else:
        
        raise ValueError('Unknown type of the order of actions!')
                                
    return Records





def RankingAlgorithm(Arguments, Records, Agent):
    
    """
    This function models how the ranking algorithm works
    """
    
    if Agent < Arguments['N_Bots']:  # this agent is a bot
        
        #N_Comments_Display = Records['Comments'].shape[0]  # the bot observes all the comments and likes (the order is not important here)
        
        # In follow-up research it would be interesting to consider the situation when bots have possibility to only a limited number of actions - for example, to avoid ban
        
        CommentIds = Records['Comments'].index
        
    else:  # this agent is not a bot - he/she has cognitive cosntrainsts
        
        N_Comments_Display = round(np.random.exponential(scale=Arguments['Scale'], size=1)[0]+1)
    
        if Arguments['RankingAlgorithmType'] == 'LikesCount':  # comments that have more likes are prioritized

            CommentIds = Records['Comments'].sort_values(by='LikesCount', ascending=False).index[:N_Comments_Display]

        elif Arguments['RankingAlgorithmType'] == 'RepliesCount': # comments that have more replies are prioritized

            CommentIds = Records['Comments'].sort_values(by='RepliesCount', ascending=False).index[:N_Comments_Display]

        elif Arguments['RankingAlgorithmType'] == 'Time': # more recent comments are prioritized

            CommentIds = Records['Comments'].sort_values(by='Time', ascending=False).index[:N_Comments_Display]

        else:

            raise ValueError('Unknown ranking algorithm type!')
    
    return (Records['Posts'].index[-1:], CommentIds)




def AgentMoves(Arguments, Records, t, Agent):
    
    """
    This function models how the ranking algorithm works
    """
    
    Records['Activations'].append(Agent)  # write down who is acting now
    
    (PostIds, CommentIds) = RankingAlgorithm(Arguments, Records, Agent)  # the agent will see these posts and comments in its news feed
    
    #------------------------------------ Proceeding posts
    
    #print(PostIds)
    
    for PostId in PostIds:
        
        Records = AgentProceedElement(Arguments, Records, t, Agent, PostId, 'Posts')
    
    #------------------------------------ Proceeding comments
    
    for CommentId in CommentIds:
        
        Records = AgentProceedElement(Arguments, Records, t, Agent, CommentId, 'Comments') 
    
    return Records




def ModelIteration(Arguments, Records, t, Type):
    
    """
    This function models a single iteration of the model
    """
    
    Agents = Records['Agents']
    
    if Type == 'Bot':
        
        Agent = np.random.choice(Agents.index[:Arguments['N_Bots']], 
                                 #p=Arguments['ActivationProbabilities'][:Arguments['N_Bots']], 
                                )  # choose a random agent (by default - unifrom distribution)
        
    elif Type == 'Native Agent':
        
        Agent = np.random.choice(Agents.index[Arguments['N_Bots']:Arguments['N']], 
                                 #p=Arguments['ActivationProbabilities'][Arguments['N_Bots'], Arguments['N']],
                                )  # choose a random agent (by default - unifrom distribution)
        
    else:
        
        raise ValueError('Unknown agent type!')
    
    Records = AgentMoves(Arguments, Records, t, Agent)  # the chosen agent does some actions
    
    return Records




################################################################################################################




def ModelRealization(Arguments):

    """
    This function models a single realization of the model
    """
    
    # Add one item to the dictionary Arguments that will store activation probabilities of agents
    
    if Arguments['ActivationProbabilitiesType'] == 'uniform':
        
        ActivationProbabilities = np.ones(Arguments['N']) #/ N
        
        #ActivationProbabilities[:Arguments['N_Bots']] = ActivationProbabilities[:Arguments['N_Bots']] / ActivationProbabilities[:Arguments['N_Bots']].sum()
        
        #ActivationProbabilities[Arguments['N_Bots']:] = ActivationProbabilities[Arguments['N_Bots']:] / ActivationProbabilities[Arguments['N_Bots']:].sum()
        
        #Arguments['ActivationProbabilities'] = ActivationProbabilities
        
    else:
        raise ValueError('Unknown type of activation probabilities!')
    
    Records = {}  # here we will store all information regarding model dynamics
    
    #------------------------------------ Initialization
    
    Agents = pd.DataFrame(columns=['Opinion(t0)'])
    
    for i in range(Arguments['N_Bots']):  # define opinions of bots
        Agents.loc[Agents.shape[0]] = [Arguments['OpinionAlphabet'][-1]]  # Bots have opinion x_m (opposite to x_1)        
    
    
    InitialFractions = (np.array(Arguments['InitialOpinionDistribution']) * (Arguments['N'] - Arguments['N_Bots'])).astype(int)
    
    # Check if we did not lose agents due to number rounding
    
    if InitialFractions.sum() != (Arguments['N'] - Arguments['N_Bots']):
        raise ValueError('Wrong number of agents!')
    else:
        pass
    
    for i in range(len(Arguments['OpinionAlphabet'])):
        
        opinion = Arguments['OpinionAlphabet'][i]
        
        for k in range(InitialFractions[i]):
            Agents.loc[Agents.shape[0]] = [opinion]            
            
    
    #for i in range(Arguments['N_Bots'], Arguments['N']):  # define opinions of native agents
    #    Agents.loc[Agents.shape[0]] = [np.random.choice(Arguments['OpinionAlphabet'], 
    #                                                    p=Arguments['InitialOpinionDistribution'])]
    
    AgentsLikes = [[] for i in range(Arguments['N'])]
    
    Fractions = [[Agents[Agents['Opinion(t0)'] == i].shape[0] for i in Arguments['OpinionAlphabet']]]
    
    #print(Fractions)
    
    Posts = pd.DataFrame(columns=['Time', 'Opinion', 'SubjectId','LikesCount', 'RepliesCount'])
    
    Comments = pd.DataFrame(columns=['Time', 'Opinion', 'SubjectId', 'ObjectId', 'ObjectType', 'LikesCount', 'RepliesCount'])
    
    Likes = pd.DataFrame(columns=['Time', 'SubjectId', 'ObjectId', 'ObjectType'])
    
    Activations = []  # here we will store activation records
    
    #------------------------------------ Add Post
    
    Posts.loc[Posts.shape[0]] = [
                                 0,                                # time creation t=0 is assumed 
                                 Arguments['OpinionAlphabet'][0],  # Post has opinion x_1
                                 -1,                               # owner of the post
                                 0,                                # no likes at the beginning
                                 0,                                # no replies at the beginning
                                ]
    
    #------------------------------------ Integrate information in the dictionary
        
    Records['Agents'] = Agents
    Records['Posts'] = Posts
    Records['Comments'] = Comments
    Records['Likes'] = Likes
    Records['Activations'] = Activations
    Records['Fractions'] = Fractions
    Records['AgentsLikes'] = AgentsLikes
    
    #------------------------------------ Model dynamics
    
    T = Arguments['T']
    
    T_bots = Arguments['T_bots']  # number of iterations when bots operate
    
    for t in range(1, T_bots+1):
        
        #print(f'Iteration {t} of {T}')
        #clear_output(wait=True)
        
        Records = ModelIteration(Arguments, Records, t, 'Bot')
        
        #Records = ModelIteration(Arguments, Records, t)
    
    #for t in range(T_bots+1, Arguments['T']+1):
    
    t = T_bots+1
    
    while True:    
        #print(f'Iteration {t} of {T}')
        #clear_output(wait=True)
        
        Records = ModelIteration(Arguments, Records, t, 'Native Agent')
        
        t = t + 1
        
        if Records['Fractions'][-1][1] == 0:  # go on while there are neutral agents (applicable only for our case)
            break
        else:
            pass
    
    Records['Fractions'] = np.array(Records['Fractions'])
    
    #print(Records['Fractions'].shape[0])
    
    return Records  # the output of the model realization is 








